import os
from PyCmpltrtok.common import get_dir_name_ext, sep
from datasets import Dataset


def read_in_bmes(path, limit=0):
    sep(path)
    assert(isinstance(limit, int) and limit >= 0)

    xdir, xbase, xext = get_dir_name_ext(path)
    if limit:
        limit_str = f'-limit-{limit}'
    else:
        limit_str = ''
    output_path = os.path.join(xdir, xbase + f'{limit_str}.hf')
    print('output path:', output_path)

    ner_tags = []
    tokens = []

    with open(path, 'r', encoding='utf8') as fin:

        this_ner_tags = []
        this_tokens = []

        cnt = -1
        last_over = False
        while True:
            # 读取一行
            xline = fin.readline()
            if not xline:  # 返回连换行都没有，是空，那说明已经读取结束
                break  # 中断读取
            cnt += 1
            if not cnt % 50:
                print(path, cnt)
            if limit:
                if cnt >= limit:
                    break
            if '\r\n' == xline[-2:]:
                xline = xline[:-2]
            elif '\n' == xline[-1:]:
                xline = xline[:-1]
            last_over = False

            print('.', end='', flush=True)

            if not xline:
                ner_tags.append(this_ner_tags)
                tokens.append(this_tokens)
                this_ner_tags = []
                this_tokens = []
                last_over = True
                continue

            token, tag = xline.split(' ')
            this_ner_tags.append(tag)
            this_tokens.append(token)

    if not last_over:
        ner_tags.append(this_ner_tags)
        tokens.append(this_tokens)

    if 0:
        for i, this_tokens in enumerate(tokens):
            sep(i)
            this_ner_tags = ner_tags[i]
            for j, token in enumerate(this_tokens):
                tag = this_ner_tags[j]
                print(token, tag, )

    ds = Dataset.from_dict({
        'tokens': tokens,
        'ner_tags': ner_tags,
    })
    print(ds)
    ds.save_to_disk(output_path)


if '__main__' == __name__:

    if 1:
        laptop_dir = r'D:\_const\large_code\from-Dr-Sun-named_entity_recognition-master\ResumeNER'
        new_pc_dir = '/home/yunpeng/datasets/cner'

        files = [
            'dev.char.bmes',
            'test.char.bmes',
            'train.char.bmes',
            'dev_small.char.bmes',
            'test_small.char.bmes',
            'train_small.char.bmes',
        ]
        
        xdir = new_pc_dir
        bmes_paths = [os.path.join(xdir, file) for file in files]
    else:
        bmes_paths = [
            # r'D:\_const\large_code\from-Dr-Sun-named_entity_recognition-master\ResumeNER\train_small.char.bmes',
            # r'D:\_const\large_code\from-Dr-Sun-named_entity_recognition-master\ResumeNER\dev_small.char.bmes',
            # r'D:\_const\large_code\from-Dr-Sun-named_entity_recognition-master\ResumeNER\test_small.char.bmes',
            # r'D:\_const\large_code\from-Dr-Sun-named_entity_recognition-master\ResumeNER\train.char.bmes',
            # r'D:\_const\large_code\from-Dr-Sun-named_entity_recognition-master\ResumeNER\dev.char.bmes',
            # r'D:\_const\large_code\from-Dr-Sun-named_entity_recognition-master\ResumeNER\test.char.bmes',
            '/home/yunpeng/datasets/cner/dev_small.char.bmes',
        ]

    # limit = 50
    limit = 0
    for path in bmes_paths:
        read_in_bmes(path, limit=limit)
